import piexif
from PIL import Image
import argparse
import piexif
import datetime


def set_exif_maker_model(input_path, output_path, maker, model):
    """设置图片的EXIF信息中的制造商、型号，如果没有EXIF时间信息则使用系统时间"""
    try:
        # 打开图片
        img = Image.open(input_path)

        # 获取现有EXIF数据，如果没有则创建新的
        exif_dict = {}
        has_build_data = False
        if img.info.get('exif'):
            exif_data = img.info['exif']
            exif_dict = piexif.load(exif_data)

            # 检查是否有build data (时间信息)
            date_tags = [
                piexif.ExifIFD.DateTimeOriginal,
                piexif.ExifIFD.DateTimeDigitized,
                piexif.ImageIFD.DateTime
            ]
            for tag in date_tags:
                if (tag in exif_dict['Exif'] or tag in exif_dict['0th']) and exif_dict.get('Exif', {}).get(tag) or exif_dict.get('0th', {}).get(tag):
                    has_build_data = True
                    break
        else:
            # 创建新的EXIF数据结构
            exif_dict = {
                '0th': {},
                'Exif': {},
                'GPS': {},
                'Interop': {},
                '1st': {},
                'thumbnail': None
            }

        # 设置制造商和型号
        exif_dict['0th'][piexif.ImageIFD.Make] = maker.encode('utf-8')
        exif_dict['0th'][piexif.ImageIFD.Model] = model.encode('utf-8')

        # 如果没有build data，则使用系统当前时间
        if not has_build_data:
            current_time = datetime.datetime.now().strftime('%Y:%m:%d %H:%M:%S')
            print(f"未找到EXIF时间信息，使用系统当前时间: {current_time}")
            current_time_bytes = current_time.encode('utf-8')
            exif_dict['0th'][piexif.ImageIFD.DateTime] = current_time_bytes
            exif_dict['Exif'][piexif.ExifIFD.DateTimeOriginal] = current_time_bytes
            exif_dict['Exif'][piexif.ExifIFD.DateTimeDigitized] = current_time_bytes

        # 将EXIF字典转换为字节数据
        exif_bytes = piexif.dump(exif_dict)

        # 保存图片并写入EXIF数据，设置高质量保存参数
        # 对于JPEG格式，quality参数范围为1-95，95为最高质量
        # optimize=True 会进行额外的压缩优化，但可能会增加保存时间
        img.save(output_path, exif=exif_bytes, quality=95, optimize=False)
        print(f"成功设置EXIF信息并保存到: {output_path}")
        return True
    except Exception as e:
        print(f"设置EXIF信息时出错: {e}")
        return False


def main():
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='设置JPG图片的EXIF信息')
    parser.add_argument('input_path', help='输入图片路径')
    parser.add_argument('output_path', help='输出图片路径')
    parser.add_argument('--maker', default='Cannice', help='制造商名称(默认: Cannice)')
    parser.add_argument('--model', default='AI glass', help='型号名称(默认: AI glass)')

    # 解析命令行参数
    args = parser.parse_args()

    # 设置EXIF信息
    set_exif_maker_model(args.input_path, args.output_path, args.maker, args.model)


if __name__ == '__main__':
    main()